# functions related to force field charges
# assign_charges(mol, ff, eqt) can assin charges according to ppf file
# dipole(mol) can calculate dipole moment of a molecule

from kuai.ff import *
from kuai.inner import list_atom_pair12

class ChargeFunction:
    def __init__(self):
        pass
    
    def get_charges(self, term, mol, ff, eqt):
        raise NotImplementedError()
    
class FunctionATC(AtomFunction, ChargeFunction):
    def unit(self, i):
        return 'e'
     
    def get_charges(self, term, mol, ff, eqt):
        result = []
        pairs = list_atom_pair12
        for i in mol.atoms:
            item = ForceFieldKey(term, [i.type])
            item = eqt.get_equal_key(item)
            if item in ff:
                result.append(ff[item][0])
            else:
                result.append(0.0)
        return result
        
class FunctionBINC(BondFunction, ChargeFunction):
    def uniform(self, atoms, ks, indexes):
        if atoms[0] > atoms[1] and ks is not None and len(ks) > 0:
            ks[0] = -ks[0]
        BondFunction.uniform(self, atoms, None, indexes)
        
    def unit(self, i):
        return 'e'
    
    def get_charges(self, term, mol, ff, eqt):
        result = [0.0] * len(mol.atoms)
        pairs = list_atom_pair12(mol)
        for i in pairs:
            self.uniform([i[0].type, i[1].type], None, i)
            item = ForceFieldKey(term, [i[0].type, i[1].type])
            item = eqt.get_equal_key(item)
            if item in ff:
                result[mol.index(i[0])] += ff[item][0]
                result[mol.index(i[1])] -= ff[item][0]
        return result
        

def get_charges(mol, ff, eqt):
    result = [0.0] * len(mol.atoms)
    terms = ff.terms()
    for i in terms:
        try:
            func = FORCE_FIELD_FUNCTIONS[i]
            if isinstance(func, ChargeFunction):
                chargeI = func.get_charges(i, mol, ff, eqt)
                assert len(chargeI) == len(result)
                for j in range(len(chargeI)):
                    result[j] += chargeI[j]
        except KeyError:
            pass
    return result

def assign_charges(mol, ff, eqt):
    charges = get_charges(mol, ff, eqt)
    assert len(charges) == len(mol.atoms)
    for i in range(len(charges)):
        mol.atoms[i].partial_charge = charges[i]
    return mol 
        

FORCE_FIELD_FUNCTIONS['ATC'] =  FunctionATC()
FORCE_FIELD_FUNCTIONS['BINC'] =  FunctionBINC()

def dipole(mol):
    result = XYZ(0, 0, 0)
    for i in mol.atoms:
        result += i.coord * i.partial_charge
    return result
